import numpy as np
import torch
import gym
from actor_critic import Actor,Critic,Share_layer

if __name__ == "__main__":
    sl = Share_layer()
    actor = Actor(sl)
    critic = Critic(sl)

    actor.load_state_dict(state_dict=torch.load('ac.pt'))
    critic.load_state_dict(state_dict=torch.load('cr.pt'))

    env = gym.make("CartPole-v1")
    state = env.reset()
    state = np.reshape(state,(1,4))
    done = False

    while not done:
        env.render()
        pred = actor(state)
        if pred[0,0]<pred[0,1]:
            action = 1
        else:
            action = 0
        state_,reward,done,info = env.step(action)
        state_ = np.reshape(state_, (1, 4))
        state = state_

    env.close()

